import numpy as np
import os
import torch
import torch.nn as nn
from torchdiffeq import odeint_adjoint as odeint
import yaml
import gym
import time
import random
import csv
from pathlib import Path
from IWDCarDynamics import IWDCarDynamics

class GPUVectorizedCarEnv:
    def __init__(self,
        preset_name,
        n,
        dt=0.01,
        solver="euler",
        device="cuda:0",
        drivetrain="iwd",
        disturbance_param=None,
        randomize_param={},
        random_seed=None,
        initial_state=None,
        **kwargs,
    ):
        if random_seed is not None:
            torch.manual_seed(random_seed)
            torch.cuda.manual_seed_all(random_seed)
            np.random.seed(random_seed)
            random.seed(random_seed)
        self.num_states = 20
        # 基于tesla_model_3调整观测空间范围
        self.observation_space = gym.spaces.Box(low=np.array([-1000.,-1000.,-50.,-200.,-200.,-20.,-1.0,-1.0,-1.0,-0.5,-0.5,-0.5,-0.5,-10.,-10.,-10.,-50.,-50.,-50.,-50.]), high=np.array([1000.,1000.,50.,200.,200.,20.,1.0,1.0,1.0,0.5,0.5,0.5,0.5,10.,10.,10.,50.,50.,50.,50.]), shape=(20,))
        self.state_space = self.observation_space
        self.drivetrain = drivetrain
        # 基于tesla_model_3调整动作空间范围
        if drivetrain == "4wd":
            self.num_actions = 6  # delta, omega*1, f1, f2, f3, f4
            self.action_space = gym.spaces.Box(low=np.array([-0.45, 0., -2000., -2000., -2000., -2000.]), high=np.array([0.45, 20., 2000., 2000., 2000., 2000.]), shape=(6,))
            self.cast_action = lambda u: torch.cat(list(map(lambda v: torch.unsqueeze(v, 1), [u[:, 0], u[:, 1], u[:, 1], u[:, 1], u[:, 1], u[:, 2], u[:, 3], u[:, 4], u[:, 5]])), 1)
        elif drivetrain == "2iwd":
            self.num_actions = 7  # delta, omega_f, omega_r, f1, f2, f3, f4
            self.action_space = gym.spaces.Box(low=np.array([-0.45, 0., 0., -2000., -2000., -2000., -2000.]), high=np.array([0.45, 20., 20., 2000., 2000., 2000., 2000.]), shape=(7,))
            self.cast_action = lambda u: torch.cat(list(map(lambda v: torch.unsqueeze(v, 1), [u[:, 0], u[:, 1], u[:, 1], u[:, 2], u[:, 2], u[:, 3], u[:, 4], u[:, 5], u[:, 6]])), 1)
        elif drivetrain == "iwd":
            self.num_actions = 9  # delta, omega_fr, omega_fl, omega_rr, omega_rl, f1, f2, f3, f4
            self.action_space = gym.spaces.Box(low=np.array([-0.45, 0., 0., 0., 0., -2000., -2000., -2000., -2000.]), high=np.array([0.45, 20., 20., 20., 20., 2000., 2000., 2000., 2000.]), shape=(9,))
            self.cast_action = lambda u: u

        self.preset_name = preset_name
        self.n = n
        self.dt = dt
        self.solver = solver
        self.device = torch.device(device)
        file_path = os.path.dirname(__file__)
        with open(os.path.join(file_path, "presets.yaml"),encoding='utf-8') as f:
            presets = yaml.safe_load(f)
            params = presets[preset_name]["parameters"]
        self.p_body = torch.zeros((n, 27), device=self.device)
        self.p_body[:, 0] = params["lF"]
        self.p_body[:, 1] = params["lR"]
        self.p_body[:, 2] = params["m"]
        self.p_body[:, 3] = params["h"]
        self.p_body[:, 4] = params["g"]
        self.p_body[:, 5] = params["Iz"]
        self.p_body[:, 6] = params["T"]

        self.p_body[:, 7] = params["wheel_radius"]
        self.p_body[:, 8] = params["m_s"]
        self.p_body[:, 9] = params["m_uf"]
        self.p_body[:, 10] = params["m_ur"]
        self.p_body[:, 11] = params["Ix"]
        self.p_body[:, 12] = params["Iy"]
        self.p_body[:, 13] = params["C_af"]
        self.p_body[:, 14] = params["C_ar"]
        self.p_body[:, 15] = params["K21"]
        self.p_body[:, 16] = params["C21"]
        self.p_body[:, 17] = params["K23"]
        self.p_body[:, 18] = params["C23"]
        self.p_body[:, 19] = params["K1"]
        self.p_body[:, 20] = params["C1"]
        self.p_body[:, 21] = params["Kaf"]
        self.p_body[:, 22] = params["Kar"]
        self.p_body[:, 23] = params["k_phi"]
        self.p_body[:, 24] = params["c_phi"]
        self.p_body[:, 25] = params["k_theta"]
        self.p_body[:, 26] = params["c_theta"]

        self.p_tyre = torch.zeros((n, 4), device=self.device)
        self.p_tyre[:, 0] = params["B"]
        self.p_tyre[:, 1] = params["C"]
        self.p_tyre[:, 2] = params["D"]
        self.p_tyre[:, 3] = params["E"]
        self.randomize_param = randomize_param
        self.initial_state = initial_state
        self.s = self.initial_state if initial_state is not None else torch.zeros((n, 20), device=self.device)
        self.current_obs = None
        self.dynamics = None
        self.disturbance_param = disturbance_param
        if disturbance_param is not None:
            # 扰动维度扩展：[fx_1..4, fy_1..4, M_phi, M_theta] = 10维
            # 前8维：轮胎力扰动 (N)
            # 后2维：力矩扰动 (N·m)，分别作用于侧倾(ddphi)和俯仰(ddtheta)
            self.disturbance = torch.zeros((self.n, 10), device=self.device)
        self.step_count = torch.zeros(self.n, dtype=torch.int64, device=self.device)
        self.total_step_count = 0
        # 为每个并行环境单独保存一条轨迹：list(n_env) -> list(steps) -> (s,u,es,obs)
        self.saved_data = [[] for _ in range(self.n)]
        self.train = kwargs.get("train", False)
        # 记录实验名称（用于按实验目录保存数据），例如：a2c_continuous_moose_default_1124_2055
        self.full_experiment_name = kwargs.get("full_experiment_name", None)
        # 异常处理配置
        self.abnormal_handling = kwargs.get("abnormal_handling", "reset_and_mark_done")  # "reset_and_mark_done", "clamp_only", "disabled"
        
        # ============ RARL 对抗学习相关 ============
        # 对抗扰动动作空间（10维）：[fx_1..4, fy_1..4, M_phi, M_theta]
        self.adv_action_space = gym.spaces.Box(
            low=np.array([-150.] * 8 + [-60000., -5000.], dtype=np.float32),
            high=np.array([150.] * 8 + [60000., 5000.], dtype=np.float32),
            dtype=np.float32
        )
        # adversary扰动缓存（由外部adversary策略设置）
        self.adversary_disturbance = None
        # RARL模式标志
        self.rarl_mode = kwargs.get("rarl_mode", False)

    def randomize_item_(self, env_mask, override, key, target):
        num = int(torch.sum(env_mask).item())
        if key in override:
            target[env_mask] = override[key]
        elif key in self.randomize_param:
            lo, hi = self.randomize_param[key]
            target[env_mask] = lo + (hi - lo) * torch.rand(num, device=self.device)

    def randomize_items_(self, env_mask, override):
        rand_item_ = lambda key, target: self.randomize_item_(env_mask, override, key, target)
        
        # 轮胎 Pacejka 参数随机化
        rand_item_("B", self.p_tyre[:, 0])
        rand_item_("C", self.p_tyre[:, 1])
        rand_item_("D", self.p_tyre[:, 2])
        rand_item_("E", self.p_tyre[:, 3])
        
        # 车身参数随机化（按照p_body数组的索引顺序）
        rand_item_("lF", self.p_body[:, 0])
        rand_item_("lR", self.p_body[:, 1])
        rand_item_("m", self.p_body[:, 2])
        rand_item_("h", self.p_body[:, 3])
        # g (重力加速度) 通常不随机化
        rand_item_("Iz", self.p_body[:, 5])
        rand_item_("T", self.p_body[:, 6])
        rand_item_("wheel_radius", self.p_body[:, 7])
        rand_item_("m_s", self.p_body[:, 8])
        rand_item_("m_uf", self.p_body[:, 9])
        rand_item_("m_ur", self.p_body[:, 10])
        rand_item_("Ix", self.p_body[:, 11])
        rand_item_("Iy", self.p_body[:, 12])
        rand_item_("C_af", self.p_body[:, 13])
        rand_item_("C_ar", self.p_body[:, 14])
        rand_item_("K21", self.p_body[:, 15])
        rand_item_("C21", self.p_body[:, 16])
        rand_item_("K23", self.p_body[:, 17])
        rand_item_("C23", self.p_body[:, 18])
        rand_item_("K1", self.p_body[:, 19])
        rand_item_("C1", self.p_body[:, 20])
        rand_item_("Kaf", self.p_body[:, 21])
        rand_item_("Kar", self.p_body[:, 22])
        rand_item_("k_phi", self.p_body[:, 23])
        rand_item_("c_phi", self.p_body[:, 24])
        rand_item_("k_theta", self.p_body[:, 25])
        rand_item_("c_theta", self.p_body[:, 26])

    def randomize(self, env_mask=None, override={}):
        if env_mask is None:
            env_mask = torch.ones(self.n, dtype=torch.bool, device=self.device)
        self.randomize_items_(env_mask, override)

    def clamp_physics_state(self):
        """对物理状态进行裁剪，防止仿真发散 - 基于tesla_model_3调整范围"""
        # 位置和角度限制（基于Tesla Model 3的实际运行范围）
        self.s[:, 0] = torch.clamp(self.s[:, 0], -1000, 1000)    # x位置 (扩大到±1km)
        self.s[:, 1] = torch.clamp(self.s[:, 1], -1000, 1000)    # y位置 (扩大到±1km)
        self.s[:, 2] = torch.clamp(self.s[:, 2], -500, 500)        # psi角度 (保持±50rad)
        
        # 速度限制（基于Tesla Model 3的性能）
        self.s[:, 3] = torch.clamp(self.s[:, 3], -200, 200)      # x速度 (扩大到±200m/s, ~720km/h)
        self.s[:, 4] = torch.clamp(self.s[:, 4], -200, 200)      # y速度 (扩大到±200m/s)
        self.s[:, 5] = torch.clamp(self.s[:, 5], -200, 200)        # psi角速度 (±20rad/s)
        
        # 悬挂状态限制（基于Tesla Model 3的悬挂行程）
        self.s[:, 6] = torch.clamp(self.s[:, 6], -4.6, 4.6)      # Zs: ±1m垂直位移
        self.s[:, 7] = torch.clamp(self.s[:, 7], -10.2, 11.2)      # phi: ±57度侧倾角
        self.s[:, 8] = torch.clamp(self.s[:, 8], -11.2, 11.2)      # theta: ±57度俯仰角
        self.s[:, 9:13] = torch.clamp(self.s[:, 9:13], -0.4, 0.4)  # Z11-Z14: ±50cm轮胎位移
        
        # 悬挂速度限制（基于Tesla Model 3的动态性能）
        self.s[:, 13] = torch.clamp(self.s[:, 13], -100, 100)      # dZs: ±10m/s垂直速度
        self.s[:, 14] = torch.clamp(self.s[:, 14], -100, 100)      # dphi: ±10rad/s侧倾角速度
        self.s[:, 15] = torch.clamp(self.s[:, 15], -100, 100)      # dtheta: ±10rad/s俯仰角速度
        self.s[:, 16:20] = torch.clamp(self.s[:, 16:20], -500, 500)  # dZ11-dZ14: ±50m/s轮胎速度
        
        # 检测异常状态并处理
        if self.abnormal_handling != "disabled":
            nan_mask = torch.isnan(self.s).any(dim=1)
            inf_mask = torch.isinf(self.s).any(dim=1)
            abnormal_mask = nan_mask | inf_mask
            
            if abnormal_mask.any():
                if self.abnormal_handling == "reset_and_mark_done":
                    print(f"Warning: {abnormal_mask.sum()} environments have abnormal states, resetting...")
                    self.reset_abnormal_envs(abnormal_mask)
                elif self.abnormal_handling == "clamp_only":
                    print(f"Warning: {abnormal_mask.sum()} environments have abnormal states, clamping...")
                    # 只裁剪，不重置

    def reset_abnormal_envs(self, abnormal_mask):
        """重置异常状态的环境"""
        num_abnormal = abnormal_mask.sum().item()
        if num_abnormal > 0:
            # 方案1：重置为初始状态并标记为done (推荐)
            if hasattr(self, 'is_done'):
                # 对于有done标记的环境(如FixedCircleIWDEnv)，标记为失败
                self.is_done[abnormal_mask] = 1  # 1表示失败
            
            # 重置异常环境的状态为初始状态
            if self.initial_state is not None:
                self.s[abnormal_mask] = self.initial_state
            else:
                self.s[abnormal_mask] = torch.zeros((num_abnormal, 20), device=self.device)
            
            # 重置步数计数
            self.step_count[abnormal_mask] = 0

    def obs(self):
        obs = self.s
        # 裁剪观测值到定义的范围内
        obs_low = torch.tensor(self.observation_space.low, device=self.device, dtype=obs.dtype)
        obs_high = torch.tensor(self.observation_space.high, device=self.device, dtype=obs.dtype)
        obs = torch.clamp(obs, obs_low, obs_high)
        return obs

    def reward(self):
        return torch.zeros(self.n, device=self.device)

    def done(self):
        return torch.zeros(self.n, device=self.device)

    def info(self):
        return {}

    def get_number_of_agents(self):
        return self.n

    def get_num_parallel(self):
        return self.n  # 这里 self.n 就是并行环境数+++-+

    def disturbed_dynamics(self):
        """
        生成扰动动力学模型，包含力扰动和力矩扰动
        
        扰动来源有两种模式：
        1. 随机AR(1)扰动（传统模式，disturbance_param != None）
        2. RARL对抗扰动（adversary_disturbance != None）
        
        扰动参数设计（参考 tesla_model_3 参数）：
        - 力扰动 (w_force = 25 N)：作用于轮胎纵向力fx和侧向力fy
        - 侧倾力矩扰动 (w_phi = 100 N·m)：参考 Ix=500 kg·m²，k_phi=13000 N·m/rad
          产生角加速度扰动 ~0.2 rad/s²，约为典型弹性力矩的 1-2%
        - 俯仰力矩扰动 (w_theta = 150 N·m)：参考 Iy=1300 kg·m²，k_theta=10000 N·m/rad
          产生角加速度扰动 ~0.12 rad/s²，约为典型弹性力矩的 1.5%
        """
        # ============ RARL模式：使用adversary网络输出的扰动 ============
        if self.rarl_mode and self.adversary_disturbance is not None:
            # adversary_disturbance 是由外部adversary策略根据当前state生成的
            # 形状: (n, 10) = [fx_1..4, fy_1..4, M_phi, M_theta]
            disturbance = self.adversary_disturbance
            return IWDCarDynamics(self.cast_action(self.u), self.p_body, self.p_tyre, disturbance=disturbance)
        
        # ============ 传统模式：AR(1)随机扰动 ============
        a, w_force = self.disturbance_param
        
        # 力矩扰动强度设计（基于车辆惯量参数）
        # 侧倾力矩扰动：考虑 Ix=500, k_phi=13000，设计 w_phi 使扰动约为弹性力矩的1-2%
        w_phi = 1500.0     # N·m，侧倾力矩扰动强度
        # 俯仰力矩扰动：考虑 Iy=1300, k_theta=10000，俯仰惯量更大，适当增加扰动
        w_theta = 750.0   # N·m，俯仰力矩扰动强度
        
        # 构建扰动强度向量 [fx_1..4, fy_1..4, M_phi, M_theta]
        w_vector = torch.tensor(
            [w_force] * 8 + [w_phi, w_theta], 
            device=self.device, 
            dtype=self.disturbance.dtype
        )
        
        # AR(1) 扰动更新
        noise = torch.randn((self.n, 10), device=self.device)
        self.disturbance = a * self.disturbance + w_vector * noise
        
        return IWDCarDynamics(self.cast_action(self.u), self.p_body, self.p_tyre, disturbance=self.disturbance)
    
    # ============ RARL 对抗学习接口方法 ============
    
    def set_adversary_disturbance(self, disturbance):
        """
        设置由adversary策略生成的扰动
        
        Args:
            disturbance: Tensor, shape (n, 10)
                扰动向量 [fx_1..4, fy_1..4, M_phi, M_theta]
                - fx_1..4: 四个轮胎的纵向力扰动 (N)
                - fy_1..4: 四个轮胎的侧向力扰动 (N)  
                - M_phi: 侧倾力矩扰动 (N·m)
                - M_theta: 俯仰力矩扰动 (N·m)
        """
        if isinstance(disturbance, np.ndarray):
            disturbance = torch.tensor(disturbance, device=self.device, dtype=torch.float32)
        self.adversary_disturbance = disturbance
    
    def clear_adversary_disturbance(self):
        """清除adversary扰动"""
        self.adversary_disturbance = None
    
    def enable_rarl_mode(self, enabled=True):
        """启用/禁用RARL对抗模式"""
        self.rarl_mode = enabled
        if not enabled:
            self.clear_adversary_disturbance()
    
    def get_adversary_action_space(self):
        """获取adversary的动作空间（扰动空间）"""
        return self.adv_action_space

    def reset(self):
        self.randomize()
        self.s = self.initial_state if self.initial_state is not None else torch.zeros((self.n, 20), device=self.device)
        self.u = torch.zeros((self.n, self.num_actions), device=self.device)
        self.dynamics = IWDCarDynamics(self.cast_action(self.u), self.p_body, self.p_tyre)
        self.es = self.dynamics.compute_extended_state(self.s)
        self.step_count = torch.zeros(self.n, dtype=torch.int64, device=self.device)
        self.total_step_count = 0
        # 重置每个并行环境的采样数据
        self.saved_data = [[] for _ in range(self.n)]
        obs = self.obs()
        self.current_obs = obs
        return obs

    def step(self, u, override_s=None):
        self.u = u

        self.dynamics = IWDCarDynamics(self.cast_action(u), self.p_body, self.p_tyre) if self.disturbance_param is None else self.disturbed_dynamics()
        if override_s is None:
            self.s = odeint(self.dynamics, self.s, torch.tensor([0., self.dt], device=self.device), method=self.solver)[1, :, :]
        else:
            self.s[:] = torch.tensor(override_s).unsqueeze(0)
        
        # 对物理状态进行裁剪，防止仿真发散
        self.clamp_physics_state()
        
        self.es = self.dynamics.compute_extended_state(self.s)
        self.step_count += 1
        self.total_step_count += 1
        obs = self.obs()
        if self.train:
            reward = self.reward()
        else:
            reward = torch.zeros(self.n, device=self.device)
        done = self.done()
        info = self.info()
        self.current_obs = obs
        # 仅在测试模式下，根据结束信号保存一次完整轨迹：
        # - 训练模式 (self.train=True)：不在这里中断，以免影响训练循环
        # - 测试模式 (self.train=False)：任一并行环境结束即可触发保存
        if (not self.train) and torch.any(done) and self.saved_data:
            basename = time.strftime("%Y%m%d-%H%M%S")
            # 确定保存目录：
            # - 如果有 full_experiment_name（通过 run.py / sacrun.py 传入），则保存到 runs/{full_experiment_name}/rollouts/{basename}
            # - 否则回退到原来的 data/{basename} 目录
            if getattr(self, "full_experiment_name", None):
                base_dir = os.path.join("runs", self.full_experiment_name, "rollouts", basename)
            else:
                base_dir = os.path.join("data", basename)

            Path(base_dir).mkdir(parents=True, exist_ok=True)

            # 为每个并行环境单独保存一个 csv：env{idx}.csv
            for env_idx in range(self.n):
                env_traj = self.saved_data[env_idx]
                if not env_traj:
                    continue

                csv_path = os.path.join(base_dir, f"env{env_idx}.csv")
                with open(csv_path, 'w') as f:
                    writer = csv.writer(f, delimiter=',')
                    for step_idx in range(len(env_traj)):
                        s = env_traj[step_idx][0]
                        u = self.cast_action(env_traj[step_idx][1].unsqueeze(0)).squeeze(0)
                        es = env_traj[step_idx][2]
                        obs = env_traj[step_idx][3]  # 使用裁剪后的观测值
                        
                        # 裁剪物理状态值到合理范围（防止CSV中出现极值）
                        def clamp_value(val, min_val=-15000, max_val=15000):
                            return max(min_val, min(max_val, val))
                        
                        x, y, psi = clamp_value(s[0].item()), clamp_value(s[1].item()), clamp_value(s[2].item())
                        x_dot, y_dot, psi_dot = clamp_value(s[3].item()), clamp_value(s[4].item()), clamp_value(s[5].item())
                        Zs, phi, theta = clamp_value(s[6].item()), clamp_value(s[7].item()), clamp_value(s[8].item())
                        Z11, Z12, Z13, Z14 = clamp_value(s[9].item()), clamp_value(s[10].item()), clamp_value(s[11].item()), clamp_value(s[12].item())
                        dZs, dphi, dtheta = clamp_value(s[13].item()), clamp_value(s[14].item()), clamp_value(s[15].item())
                        dZ11, dZ12, dZ13, dZ14 = clamp_value(s[16].item()), clamp_value(s[17].item()), clamp_value(s[18].item()), clamp_value(s[19].item())
                        
                        delta, omega_fr, omega_fl, omega_rr, omega_rl = clamp_value(u[0].item()), clamp_value(u[1].item()), clamp_value(u[2].item()), clamp_value(u[3].item()), clamp_value(u[4].item())
                        f1, f2, f3, f4 = clamp_value(u[5].item()), clamp_value(u[6].item()), clamp_value(u[7].item()), clamp_value(u[8].item())
                        r = clamp_value(es[2].item())
                        beta = clamp_value(es[21].item())
                        v = clamp_value(es[20].item())
                        
                        # 侧翻风险指标
                        LLTR_front = clamp_value(es[70].item())
                        LLTR_rear = clamp_value(es[71].item())
                        gamma_fr = clamp_value(es[72].item())
                        gamma_fl = clamp_value(es[73].item())
                        gamma_rr = clamp_value(es[74].item())
                        gamma_rl = clamp_value(es[75].item())
                        
                        writer.writerow([x, y, psi, x_dot, y_dot, psi_dot, delta, omega_fr, omega_fl, omega_rr, omega_rl, f1, f2, f3, f4, r, beta, v, Zs, phi, theta, Z11, Z12, Z13, Z14, dZs, dphi, dtheta, dZ11, dZ12, dZ13, dZ14, LLTR_front, LLTR_rear, gamma_fr, gamma_fl, gamma_rr, gamma_rl])

            torch.save(self.saved_data, os.path.join(base_dir, "rollout.pth"))
            print("Total steps:", self.total_step_count)
            print(f"Saved rollout data for {self.n} environments to directory: {base_dir}")
            exit(0)
        return obs, reward, done, info

    def render(self, **kwargs):
        """Save rollout data to emulate rendering."""
        # 保存所有并行环境在当前步的状态，方便分别导出 csv
        for env_idx in range(self.n):
            self.saved_data[env_idx].append(
                (self.s[env_idx, :].cpu(),
                 self.u[env_idx, :].cpu(),
                 self.es[env_idx, :].cpu(),
                 self.current_obs[env_idx, :].cpu())
            )

    def detach(self):
        """Clear the gradient stored in the current state of the environment."""
        self.s = self.s.detach()
        self.u = self.u.detach()
        self.es = self.es.detach()

if __name__ == "__main__":
    from matplotlib import pyplot as plt
    from tqdm import tqdm
    import time

    # Compare solvers
    env_euler = GPUVectorizedCarEnv("xcar", 1, solver="euler", drivetrain="iwd")
    env_rk = GPUVectorizedCarEnv("xcar", 1, solver="dopri5", drivetrain="iwd")
    
    traj_euler = [env_euler.reset().cpu().numpy()]
    traj_rk = [env_rk.reset().cpu().numpy()]
    
    # Test trajectory with varying inputs
    for i in range(500):
        if i < 100:
            u = [0., 2., 2., 2., 2.]  # Straight acceleration
        elif i < 200:
            u = [0.4, 4., 3., 4., 3.]  # Left turn with differential speeds
        elif i < 300:
            u = [-0.4, 3., 4., 3., 4.]  # Right turn with differential speeds
        elif i < 400:
            u = [0., 4., -4., 4., -4.]  # Spin in place
        else:
            u = [0.2, 3., 3., 3., 3.]  # Gentle right turn
            
        s_euler, _, _, _ = env_euler.step(torch.tensor([u], device=torch.device("cuda:0")))
        s_rk, _, _, _ = env_rk.step(torch.tensor([u], device=torch.device("cuda:0")))
        traj_euler.append(s_euler.cpu().numpy())
        traj_rk.append(s_rk.cpu().numpy())
        print(f"Step {i} done")

    # Plot trajectories
    plt.figure(dpi=300)
    plt.plot([s[0][0] for s in traj_euler], [s[0][1] for s in traj_euler], label="Euler")
    plt.plot([s[0][0] for s in traj_rk], [s[0][1] for s in traj_rk], label="RK5")
    plt.legend()
    plt.axis("equal")
    plt.title("Half-DOF Car Model Trajectories")
    plt.xlabel("X position")
    plt.ylabel("Y position")
    plt.savefig("half_dof_car_trajectories.png")

    plt.figure(dpi=300)
    plt.plot([s[0][2] for s in traj_euler], label="Euler")
    plt.plot([s[0][0] for s in traj_rk], [s[0][1] for s in traj_rk], label="RK5")
    plt.legend()
    plt.title("Half-DOF Car Model Trajectories")
    plt.xlabel("Time step")
    plt.ylabel("Psi")
    plt.savefig("half_dof_car_psi.png")
    